#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Output generation component
"""

import json
import pandas as pd
from typing import Dict, Any, List, Optional
from data_engine.models.data_models import Dataset, ChartAnalysis, ProcessingResult, VisualizationResult
from data_engine.config import ConfigManager, ChartType
from data_engine.utils.logger import engine_logger


class OutputGenerator:
    """输出生成器"""
    
    def __init__(self, llm_client):
        """
        初始化输出生成器
        
        Args:
            llm_client: LiteLLM客户端实例
        """
        self.llm_client = llm_client
        self.config_manager = ConfigManager()
    
    async def generate_visualization(self, dataset: Dataset, analysis: ChartAnalysis, 
                                   processing_result: ProcessingResult) -> VisualizationResult:
        """
        生成可视化配置
        
        Args:
            dataset: 原始数据集
            analysis: 图表分析结果
            processing_result: 数据处理结果
            
        Returns:
            VisualizationResult: 可视化结果
        """
        start_time = pd.Timestamp.now()
        
        try:
            if analysis.chart_type == "table" or not analysis.need_chart:
                # 生成表格配置
                table_config = await self._generate_table_config_with_llm(dataset, processing_result)
                # 确保表格配置符合标准格式
                config = self._format_table_output(table_config)
                result_type = "table"
            else:
                # 生成ECharts配置
                echarts_option = await self._generate_echarts_config_with_llm(dataset, analysis, processing_result)
                # 确保ECharts配置包装在option字段中
                config = self._format_echarts_output(echarts_option)
                result_type = "echarts"
            
            generation_time = (pd.Timestamp.now() - start_time).total_seconds()
            
            return VisualizationResult(
                result_type=result_type,
                config=config,
                chart_analysis=analysis,
                processing_result=processing_result,
                generation_time=generation_time,
                success=True
            )
            
        except Exception as e:
            generation_time = (pd.Timestamp.now() - start_time).total_seconds()
            engine_logger.error(f"可视化生成失败: {str(e)}")
            
            # 生成备用表格配置
            fallback_table = self._generate_fallback_table(processing_result.processed_df)
            fallback_config = self._format_table_output(fallback_table)
            
            return VisualizationResult(
                result_type="table",
                config=fallback_config,
                chart_analysis=analysis,
                processing_result=processing_result,
                generation_time=generation_time,
                success=False,
                error_message=str(e)
            )
    
    async def _generate_echarts_config_with_llm(self, dataset: Dataset, analysis: ChartAnalysis, 
                                               processing_result: ProcessingResult) -> Dict[str, Any]:
        """使用LLM生成ECharts配置"""
        engine_logger.info(f"使用LLM生成{analysis.chart_type}图表配置")
        
        # 准备数据
        df = processing_result.processed_df
        
        # 限制数据样本大小以避免prompt过长
        sample_size = min(10, len(df))
        sample_df = df.head(sample_size)
        
        # 构建字段映射信息
        field_mapping = {
            "chart_type": analysis.chart_type,
            "x_field": analysis.x_field,
            "y_field": analysis.y_field,
            "y_fields": analysis.y_fields,
            "name_field": analysis.name_field,
            "value_field": analysis.value_field
        }
        
        # 智能构建字段描述信息
        field_descriptions = {}
        for col in df.columns:
            if (dataset.field_descriptions and 
                col in dataset.field_descriptions and 
                dataset.field_descriptions[col]):
                field_descriptions[col] = dataset.field_descriptions[col]
            else:
                # 如果没有字段描述，使用字段名和数据类型
                dtype_info = str(df[col].dtype)
                unique_count = df[col].nunique()
                field_descriptions[col] = f"{col} ({dtype_info}, {unique_count} unique values)"
        
        # 智能组装数据集描述信息
        dataset_info_parts = []
        if dataset.dataset_description:
            dataset_info_parts.append(f"数据集描述: {dataset.dataset_description}")
        if dataset.table_description:
            dataset_info_parts.append(f"表格描述: {dataset.table_description}")
        
        dataset_info = "\n".join(dataset_info_parts) if dataset_info_parts else "请根据数据字段和样本数据理解数据含义"
        
        # 获取提示词模板并构建提示词
        template = self.config_manager.get_prompt_template("echarts_generation")
        prompt = template.format(
            chart_type=analysis.chart_type,
            processed_data=sample_df.to_json(orient='records', indent=2),
            field_mapping=json.dumps(field_mapping, indent=2),
            field_descriptions=json.dumps(field_descriptions, indent=2, ensure_ascii=False),
            dataset_info=dataset_info
        )
        
        try:
            # 调用LLM生成配置
            response = await self.llm_client.acompletion(
                messages=[{"role": "user", "content": prompt}]
            )
            
            result_text = self.llm_client.get_response_content(response).strip()
            
            # 解析JSON配置
            if result_text.startswith('```json'):
                result_text = result_text.replace('```json', '').replace('```', '').strip()
            elif result_text.startswith('```'):
                result_text = result_text.replace('```', '').strip()
            
            config = json.loads(result_text)
            
            # 验证和修复配置
            config = self._fix_echarts_config(config, df, analysis)
            
            return config
            
        except Exception as e:
            engine_logger.warning(f"LLM生成配置失败: {str(e)}，使用备用方法")
            return self._generate_echarts_config_fallback(df, analysis, dataset)
    
    def _fix_echarts_config(self, config: Dict[str, Any], df: pd.DataFrame, 
                           analysis: ChartAnalysis) -> Dict[str, Any]:
        """修复和完善ECharts配置"""
        
        # 确保有series配置
        if "series" not in config or not config["series"]:
            config["series"] = []
        
        # 处理数据映射
        if analysis.chart_type in ["line", "bar", "area"]:
            # 确保有x轴数据
            if "xAxis" in config and analysis.x_field in df.columns:
                if df[analysis.x_field].dtype == 'object':
                    config["xAxis"]["data"] = df[analysis.x_field].tolist()
                
            # 修复series数据
            for i, series in enumerate(config["series"]):
                if analysis.y_fields and i < len(analysis.y_fields):
                    y_field = analysis.y_fields[i]
                elif analysis.y_field:
                    y_field = analysis.y_field
                else:
                    continue
                
                if y_field in df.columns:
                    if df[analysis.x_field].dtype == 'object':
                        # 类别型x轴，只需要y值数组
                        series["data"] = df[y_field].tolist()
                    else:
                        # 数值型x轴，需要[x,y]坐标对
                        series["data"] = [[x, y] for x, y in zip(df[analysis.x_field], df[y_field])]
        
        elif analysis.chart_type == "pie":
            if config["series"] and analysis.name_field in df.columns and analysis.value_field in df.columns:
                pie_data = []
                for _, row in df.iterrows():
                    pie_data.append({
                        "name": str(row[analysis.name_field]),
                        "value": float(row[analysis.value_field])
                    })
                config["series"][0]["data"] = pie_data
        
        elif analysis.chart_type == "scatter":
            if config["series"] and analysis.x_field in df.columns and analysis.y_field in df.columns:
                scatter_data = [[float(x), float(y)] for x, y in zip(df[analysis.x_field], df[analysis.y_field])]
                config["series"][0]["data"] = scatter_data
        
        return config
    
    def _generate_echarts_config_fallback(self, df: pd.DataFrame, analysis: ChartAnalysis, 
                                         dataset: Dataset) -> Dict[str, Any]:
        """备用ECharts配置生成方法 - 返回纯option配置"""
        if analysis.chart_type == "pie":
            return self._generate_pie_config(df, analysis, dataset)
        elif analysis.chart_type in ["line", "bar", "scatter", "area"]:
            return self._generate_xy_chart_config(df, analysis, dataset)
        else:
            raise ValueError(f"不支持的图表类型: {analysis.chart_type}")
    
    async def _generate_table_config_with_llm(self, dataset: Dataset, 
                                             processing_result: ProcessingResult) -> Dict[str, Any]:
        """使用LLM生成表格配置"""
        engine_logger.info("使用LLM生成表格配置")
        
        df = processing_result.processed_df
        
        # 限制数据样本大小
        sample_size = min(10, len(df))
        sample_df = df.head(sample_size)
        
        # 智能构建字段描述信息
        field_descriptions = {}
        for col in df.columns:
            if (dataset.field_descriptions and 
                col in dataset.field_descriptions and 
                dataset.field_descriptions[col]):
                field_descriptions[col] = dataset.field_descriptions[col]
            else:
                # 如果没有字段描述，使用字段名和数据类型
                dtype_info = str(df[col].dtype)
                unique_count = df[col].nunique()
                field_descriptions[col] = f"{col} ({dtype_info}, {unique_count} unique values)"
        
        # 智能组装数据集描述信息
        dataset_info_parts = []
        if dataset.dataset_description:
            dataset_info_parts.append(f"数据集描述: {dataset.dataset_description}")
        if dataset.table_description:
            dataset_info_parts.append(f"表格描述: {dataset.table_description}")
        
        dataset_info = "\n".join(dataset_info_parts) if dataset_info_parts else "请根据数据字段和样本数据理解数据含义"
        
        # 获取提示词模板并构建提示词
        template = self.config_manager.get_prompt_template("table_generation")
        prompt = template.format(
            processed_data=sample_df.to_json(orient='records', indent=2),
            field_descriptions=json.dumps(field_descriptions, indent=2, ensure_ascii=False),
            dataset_info=dataset_info
        )
        
        try:
            # 调用LLM生成配置
            response = await self.llm_client.acompletion(
                messages=[{"role": "user", "content": prompt}]
            )
            
            result_text = self.llm_client.get_response_content(response).strip()
            
            # 解析JSON配置
            if result_text.startswith('```json'):
                result_text = result_text.replace('```json', '').replace('```', '').strip()
            elif result_text.startswith('```'):
                result_text = result_text.replace('```', '').strip()
            
            config = json.loads(result_text)
            
            # 确保数据完整性
            config["data"] = df.to_dict('records')
            config["total"] = len(df)
            
            return config
            
        except Exception as e:
            engine_logger.warning(f"LLM生成表格配置失败: {str(e)}，使用备用方法")
            return self._generate_table_config_fallback(dataset, processing_result)
    
    def _generate_pie_config(self, df: pd.DataFrame, analysis: ChartAnalysis, 
                            dataset: Dataset) -> Dict[str, Any]:
        """生成饼图配置 - 返回纯ECharts option"""
        if not analysis.name_field or not analysis.value_field:
            raise ValueError("饼图需要name_field和value_field")
        
        # 准备数据
        pie_data = []
        for _, row in df.iterrows():
            pie_data.append({
                "name": str(row[analysis.name_field]),
                "value": float(row[analysis.value_field])
            })
        
        # 智能获取字段描述和配置
        if (dataset.field_descriptions and 
            analysis.name_field in dataset.field_descriptions and 
            dataset.field_descriptions[analysis.name_field]):
            name_desc = dataset.field_descriptions[analysis.name_field]
        else:
            name_desc = analysis.name_field
            
        if (dataset.field_descriptions and 
            analysis.value_field in dataset.field_descriptions and 
            dataset.field_descriptions[analysis.value_field]):
            value_desc = dataset.field_descriptions[analysis.value_field]
        else:
            value_desc = analysis.value_field
        chart_config = self.config_manager.get_chart_config("pie")
        
        # 基础配置 - 返回纯ECharts option
        config = {
            "title": {"text": f"{name_desc}分布图"},
            "tooltip": {
                "trigger": "item",
                "formatter": "{b}: {c} ({d}%)"
            },
            "legend": {
                "orient": "vertical",
                "left": "left"
            },
            "series": [{
                "name": value_desc,
                "type": "pie",
                "data": pie_data,
                **chart_config.get("default_config", {})
            }]
        }
        
        return config
    
    def _generate_xy_chart_config(self, df: pd.DataFrame, analysis: ChartAnalysis, 
                                 dataset: Dataset) -> Dict[str, Any]:
        """生成XY轴图表配置 - 返回纯ECharts option"""
        if not analysis.x_field:
            raise ValueError("XY轴图表需要x_field")
        
        # 确定y字段列表
        y_fields = analysis.y_fields if analysis.y_fields else ([analysis.y_field] if analysis.y_field else [])
        if not y_fields:
            raise ValueError("XY轴图表需要至少一个y_field")
        
        # 准备数据
        x_data = df[analysis.x_field].unique().tolist() if df[analysis.x_field].dtype == 'object' else None
        # 智能获取x字段描述
        if (dataset.field_descriptions and 
            analysis.x_field in dataset.field_descriptions and 
            dataset.field_descriptions[analysis.x_field]):
            x_desc = dataset.field_descriptions[analysis.x_field]
        else:
            x_desc = analysis.x_field
        
        # 获取图表配置
        chart_config = self.config_manager.get_chart_config(analysis.chart_type)
        echarts_type = chart_config.get("echarts_type", "line")
        
        # 生成多序列配置
        series = []
        for y_field in y_fields:
            if y_field not in df.columns:
                continue
            
            # 智能获取y字段描述
            if (dataset.field_descriptions and 
                y_field in dataset.field_descriptions and 
                dataset.field_descriptions[y_field]):
                y_desc = dataset.field_descriptions[y_field]
            else:
                y_desc = y_field
            
            # 根据x轴类型准备数据
            if df[analysis.x_field].dtype == 'object':
                # 类别型x轴，需要按x轴分组聚合
                if analysis.group_by:
                    # 如果有分组，按x轴和分组字段聚合
                    grouped_data = []
                    for x_val in x_data:
                        x_subset = df[df[analysis.x_field] == x_val]
                        if len(x_subset) > 0:
                            if analysis.aggregation == "sum":
                                value = x_subset[y_field].sum()
                            elif analysis.aggregation == "mean":
                                value = x_subset[y_field].mean()
                            elif analysis.aggregation == "count":
                                value = len(x_subset)
                            else:
                                value = x_subset[y_field].iloc[0]
                            grouped_data.append(value)
                        else:
                            grouped_data.append(0)
                    y_data = grouped_data
                else:
                    y_data = df[y_field].tolist()
            else:
                # 数值型x轴，使用坐标对
                y_data = [[x, y] for x, y in zip(df[analysis.x_field], df[y_field])]
            
            series_config = {
                "name": y_desc,
                "type": echarts_type,
                "data": y_data,
                **chart_config.get("default_config", {})
            }
            
            series.append(series_config)
        
        # 生成标题
        if len(y_fields) == 1:
            # 智能获取y字段描述
            if (dataset.field_descriptions and 
                y_fields[0] in dataset.field_descriptions and 
                dataset.field_descriptions[y_fields[0]]):
                y_desc = dataset.field_descriptions[y_fields[0]]
            else:
                y_desc = y_fields[0]
            title = f"{x_desc} vs {y_desc}"
        else:
            y_names = []
            for y in y_fields:
                if (dataset.field_descriptions and 
                    y in dataset.field_descriptions and 
                    dataset.field_descriptions[y]):
                    y_names.append(dataset.field_descriptions[y])
                else:
                    y_names.append(y)
            title = f"{x_desc} vs {', '.join(y_names)}"
        
        # 基础配置 - 返回纯ECharts option
        config = {
            "title": {"text": title},
            "tooltip": {"trigger": "axis"},
            "legend": {"data": [s["name"] for s in series]} if len(series) > 1 else {},
            "xAxis": {
                "type": "category" if df[analysis.x_field].dtype == 'object' else "value",
                "data": x_data if df[analysis.x_field].dtype == 'object' else None,
                "name": x_desc
            },
            "yAxis": {
                "type": "value",
                "name": "值" if len(y_fields) > 1 else (
                    dataset.field_descriptions.get(y_fields[0]) 
                    if (dataset.field_descriptions and 
                        y_fields[0] in dataset.field_descriptions and 
                        dataset.field_descriptions[y_fields[0]]) 
                    else y_fields[0]
                )
            },
            "series": series
        }
        
        return config
    
    def _generate_table_config_fallback(self, dataset: Dataset, 
                                       processing_result: ProcessingResult) -> Dict[str, Any]:
        """备用表格配置生成方法"""
        engine_logger.info("使用备用方法生成表格配置")
        
        df = processing_result.processed_df
        
        # 生成列配置
        columns = []
        for col in df.columns:
            # 智能获取字段描述
            if (dataset.field_descriptions and 
                col in dataset.field_descriptions and 
                dataset.field_descriptions[col]):
                field_desc = dataset.field_descriptions[col]
            else:
                # 如果没有字段描述，使用字段名
                field_desc = col
            
            # 确定数据类型
            data_type = self._get_column_data_type(df[col])
            
            columns.append({
                "key": col,
                "title": field_desc,
                "dataType": data_type
            })
        
        # 生成表格配置
        config = {
            "type": "table",
            "columns": columns,
            "data": df.to_dict('records'),
            "pagination": True,
            "pageSize": 10,
            "total": len(df)
        }
        
        return config
    
    def _get_column_data_type(self, series: pd.Series) -> str:
        """获取列的数据类型"""
        if pd.api.types.is_numeric_dtype(series):
            return "number"
        elif pd.api.types.is_datetime64_any_dtype(series):
            return "date"
        else:
            return "string"
    
    def _generate_fallback_table(self, df: pd.DataFrame) -> Dict[str, Any]:
        """生成备用表格配置 - 返回标准表格格式"""
        # 生成列配置
        columns = []
        for col in df.columns:
            data_type = self._get_column_data_type(df[col])
            columns.append({
                "key": col,
                "title": col,  # 备用情况下使用字段名作为标题
                "dataType": data_type
            })
        
        # 生成数据
        data = []
        for _, row in df.iterrows():
            row_data = {}
            for col in df.columns:
                value = row[col]
                # 处理特殊值
                if pd.isna(value):
                    row_data[col] = None
                elif isinstance(value, (int, float)):
                    row_data[col] = float(value) if pd.notna(value) else None
                else:
                    row_data[col] = str(value)
            data.append(row_data)
        
        return {
            "type": "table",
            "columns": columns,
            "data": data,
            "pagination": True,
            "pageSize": 10,
            "total": len(data)
        }
    
    def enhance_chart_config(self, config: Dict[str, Any], chart_type: str) -> Dict[str, Any]:
        """增强图表配置"""
        enhanced_config = config.copy()
        
        # 添加通用增强
        if "tooltip" not in enhanced_config:
            enhanced_config["tooltip"] = {"trigger": "axis"}
        
        if "grid" not in enhanced_config:
            enhanced_config["grid"] = {
                "left": "3%",
                "right": "4%",
                "bottom": "3%",
                "containLabel": True
            }
        
        # 根据图表类型添加特定增强
        if chart_type == ChartType.BAR.value:
            enhanced_config["tooltip"]["trigger"] = "axis"
            enhanced_config["tooltip"]["axisPointer"] = {"type": "shadow"}
        
        elif chart_type == ChartType.LINE.value:
            enhanced_config["tooltip"]["trigger"] = "axis"
            enhanced_config["tooltip"]["axisPointer"] = {"type": "cross"}
        
        elif chart_type == ChartType.PIE.value:
            enhanced_config["tooltip"]["trigger"] = "item"
            enhanced_config["tooltip"]["formatter"] = "{b}: {c} ({d}%)"
        
        return enhanced_config
    
    def validate_chart_config(self, config: Dict[str, Any], chart_type: str) -> Dict[str, Any]:
        """验证图表配置"""
        validation_result = {
            "is_valid": True,
            "errors": [],
            "warnings": []
        }
        
        # 检查必需字段
        if "series" not in config:
            validation_result["errors"].append("缺少series配置")
            validation_result["is_valid"] = False
        
        # 检查series配置
        if "series" in config:
            series = config["series"]
            if not isinstance(series, list) or len(series) == 0:
                validation_result["errors"].append("series必须是非空列表")
                validation_result["is_valid"] = False
            else:
                for i, serie in enumerate(series):
                    if "data" not in serie:
                        validation_result["errors"].append(f"series[{i}]缺少data字段")
                        validation_result["is_valid"] = False
        
        # 根据图表类型检查特定配置
        if chart_type in [ChartType.LINE.value, ChartType.BAR.value, ChartType.SCATTER.value]:
            if "xAxis" not in config:
                validation_result["warnings"].append("建议添加xAxis配置")
            if "yAxis" not in config:
                validation_result["warnings"].append("建议添加yAxis配置")
        
        return validation_result
    
    def _format_echarts_output(self, echarts_option: Dict[str, Any]) -> Dict[str, Any]:
        """格式化ECharts输出，确保包装在option字段中"""
        return {
            "type": "echarts",
            "option": echarts_option
        }
    
    def _format_table_output(self, table_config: Dict[str, Any]) -> Dict[str, Any]:
        """格式化表格输出，确保符合标准格式"""
        # 如果已经有type字段，直接返回
        if "type" in table_config:
            return table_config
        
        # 否则包装标准格式
        return {
            "type": "table",
            "columns": table_config.get("columns", []),
            "data": table_config.get("data", []),
            "pagination": table_config.get("pagination", True),
            "pageSize": table_config.get("pageSize", 10),
            "total": table_config.get("total", 0)
        }